package goroutinepool

import (
	"context"
	"fmt"
	"math/rand"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

// Priority 定义协程优先级
type Priority int

const (
	PriorityLow Priority = iota
	PriorityNormal
	PriorityHigh
)

// TaskConfig 任务配置
type TaskConfig struct {
	CPUCore  int           // 指定CPU核心，-1表示不指定
	Priority Priority      // 优先级
	Timeout  time.Duration // 运行时间限制，0表示无限制
}

// Task 任务定义
type Task struct {
	Name   string
	Fn     func(ctx context.Context) error
	Config TaskConfig
}

// Pool 协程池
type Pool struct {
	wg             sync.WaitGroup
	running        atomic.Int64
	success        atomic.Int64
	failed         atomic.Int64
	stopMonitor    chan struct{}
	monitorStopped chan struct{}
	started        atomic.Bool
	mu             sync.Mutex
	once           sync.Once // 用于保证monitor只启动一次

	// 每个CPU核心最大协程数配置
	maxGoroutinesPerCore int
	// 每个CPU核心当前运行的协程数
	coreGoroutineCounts []atomic.Int64
	numCPU              int
	// CPU亲和性设置缓存
	cpuAffinitySupported bool
}

// NewPool 创建新的协程池（每个CPU核心无协程数限制）
func NewPool(bufferSize int) *Pool {
	numCPU := runtime.NumCPU()
	coreCounts := make([]atomic.Int64, numCPU)
	for i := 0; i < numCPU; i++ {
		coreCounts[i] = atomic.Int64{}
	}

	return &Pool{
		stopMonitor:          make(chan struct{}),
		monitorStopped:       make(chan struct{}),
		coreGoroutineCounts:  coreCounts,
		numCPU:               numCPU,
		cpuAffinitySupported: runtime.GOOS == "linux", // 目前仅在Linux上支持CPU亲和性
	}
}

// NewPoolWithMaxPerCore 创建新的协程池，支持设置每个CPU核心最大协程数
// maxPerCore: 每个CPU核心允许的最大协程数，<=0表示无限制
func NewPoolWithMaxPerCore(maxPer int, maxPerCore int) *Pool {
	numCPU := runtime.NumCPU()
	coreCounts := make([]atomic.Int64, numCPU)
	for i := 0; i < numCPU; i++ {
		coreCounts[i] = atomic.Int64{}
	}

	return &Pool{
		stopMonitor:          make(chan struct{}),
		monitorStopped:       make(chan struct{}),
		maxGoroutinesPerCore: maxPerCore,
		coreGoroutineCounts:  coreCounts,
		numCPU:               numCPU,
		cpuAffinitySupported: runtime.GOOS == "linux",
	}
}

// Start 启动协程池（启动监控）
func (p *Pool) Start() {
	p.once.Do(func() {
		p.started.Store(true)
		// 启动监控协程
		go p.monitor()
	})
}

// acquireSlot 获取可用的协程槽位
// 返回要使用的CPU核心索引，如果无法获取则返回-1
func (p *Pool) acquireSlot(task *Task) int {
	// 如果没有设置最大协程数限制，直接返回配置的CPU核心或随机选择
	if p.maxGoroutinesPerCore <= 0 {
		if task.Config.CPUCore >= 0 {
			return task.Config.CPUCore
		}
		return -1 // -1表示不绑定到特定核心
	}

	// 如果指定了CPU核心
	if task.Config.CPUCore >= 0 {
		core := task.Config.CPUCore
		if core >= 0 && core < len(p.coreGoroutineCounts) {
			// 使用CAS（Compare-And-Swap）原子操作尝试获取槽位
			for {
				current := p.coreGoroutineCounts[core].Load()
				if current >= int64(p.maxGoroutinesPerCore) {
					// 该核心已满
					return -1
				}
				// 尝试原子增加
				if p.coreGoroutineCounts[core].CompareAndSwap(current, current+1) {
					return core
				}
				// CAS失败，重试
			}
		}
		return -1
	}

	// 未指定CPU核心，选择负载最轻的核心（使用最少协程的核心）
	// 优化：使用快速随机选择代替全遍历
	// 先尝试随机选择几个核心，找到可用且负载最小的
	const maxAttempts = 4
	bestCore := -1
	bestCount := int64(-1)

	// 如果核心数较少（全遍历）
	if p.numCPU <= maxAttempts {
		for i := 0; i < p.numCPU; i++ {
			current := p.coreGoroutineCounts[i].Load()
			if current < int64(p.maxGoroutinesPerCore) {
				if bestCore == -1 || current < bestCount {
					bestCore = i
					bestCount = current
					if bestCount == 0 {
						break // 找到空闲核心，立即退出
					}
				}
			}
		}
	} else {
		// 核心数较多，随机采样
		attempts := 0
		for attempts < maxAttempts {
			coreIdx := rand.Intn(p.numCPU)
			current := p.coreGoroutineCounts[coreIdx].Load()
			if current < int64(p.maxGoroutinesPerCore) {
				if bestCore == -1 || current < bestCount {
					bestCore = coreIdx
					bestCount = current
					if bestCount == 0 {
						break // 找到空闲核心，立即退出
					}
				}
			}
			attempts++
		}

		// 如果没找到合适的，尝试全遍历
		if bestCore == -1 {
			for i := 0; i < p.numCPU; i++ {
				current := p.coreGoroutineCounts[i].Load()
				if current < int64(p.maxGoroutinesPerCore) {
					if bestCore == -1 || current < bestCount {
						bestCore = i
						bestCount = current
						if bestCount == 0 {
							break
						}
					}
				}
			}
		}
	}

	if bestCore >= 0 {
		p.coreGoroutineCounts[bestCore].Add(1)
	}

	return bestCore
}

// releaseSlot 释放协程槽位
func (p *Pool) releaseSlot(core int) {
	if p.maxGoroutinesPerCore <= 0 || core < 0 {
		return
	}

	if core >= 0 && core < len(p.coreGoroutineCounts) {
		// 使用原子操作减少计数
		p.coreGoroutineCounts[core].Add(-1)
	}
}

// Submit 提交单个任务
func (p *Pool) Submit(task *Task) {
	p.submitTask(task)
}

// SubmitMultiple 批量提交多个任务
func (p *Pool) SubmitMultiple(tasks []*Task) {
	for _, task := range tasks {
		p.submitTask(task)
	}
}

// submitTask 内部任务提交方法
func (p *Pool) submitTask(task *Task) {
	if !p.started.Load() {
		p.Start()
	}

	// 尝试获取协程槽位
	core := p.acquireSlot(task)
	if core >= 0 || p.maxGoroutinesPerCore <= 0 {
		// 获取成功或无限制，启动任务
		p.wg.Add(1)
		p.running.Add(1)

		go p.runTask(task, core)
	} else {
		// 无法获取槽位，任务被拒绝
		p.failed.Add(1)
		// 减少日志输出频率以提高性能
		if p.failed.Load()%10 == 0 { // 每10个失败任务输出一次警告
			fmt.Printf("[WARNING] Task '%s' rejected: all CPU cores are at maximum capacity (%d/%d)\n",
				task.Name, p.maxGoroutinesPerCore, p.maxGoroutinesPerCore)
		}
	}
}

// Wait 等待所有协程完成
func (p *Pool) Wait() {
	p.wg.Wait()

	// 停止监控
	if p.started.Load() {
		close(p.stopMonitor)
		<-p.monitorStopped
		p.started.Store(false)
	}
}

// GetRunningCount 获取当前运行中的协程数量
func (p *Pool) GetRunningCount() int64 {
	return p.running.Load()
}

// GetSuccessCount 获取成功完成的任务数量
func (p *Pool) GetSuccessCount() int64 {
	return p.success.Load()
}

// GetFailedCount 获取失败的任务数量
func (p *Pool) GetFailedCount() int64 {
	return p.failed.Load()
}

// GetTotalCompletedCount 获取总完成数量（成功+失败）
func (p *Pool) GetTotalCompletedCount() int64 {
	return p.success.Load() + p.failed.Load()
}

// runTask 执行任务
func (p *Pool) runTask(task *Task, core int) {
	defer func() {
		p.running.Add(-1)
		p.wg.Done()

		// 释放协程槽位
		if p.maxGoroutinesPerCore > 0 {
			p.releaseSlot(core)
		}

		if r := recover(); r != nil {
			p.failed.Add(1)
			// 减少panic输出频率
			if p.failed.Load()%100 == 0 { // 每100个失败任务输出一次错误
				fmt.Printf("[ERROR] Task '%s' panicked: %v\n", task.Name, r)
			}
		}
	}()

	// 根据优先级处理（高优先级立即执行，低优先级稍微延迟）
	switch task.Config.Priority {
	case PriorityLow:
		time.Sleep(5 * time.Millisecond) // 减少延迟时间
	case PriorityHigh:
		// 立即执行
	case PriorityNormal:
		// 正常执行
	}

	// 设置CPU亲和性（仅在支持且有指定核心时）
	effectiveCore := core
	if task.Config.CPUCore >= 0 {
		effectiveCore = task.Config.CPUCore
	}
	// 优化：只有在明确指定且支持时才设置CPU亲和性
	if p.cpuAffinitySupported && effectiveCore >= 0 && effectiveCore < p.numCPU {
		if err := setCPUAffinity(effectiveCore); err == nil {
			defer runtime.UnlockOSThread()
		}
	}

	// 创建context
	ctx := context.Background()
	var cancel context.CancelFunc

	if task.Config.Timeout > 0 {
		ctx, cancel = context.WithTimeout(ctx, task.Config.Timeout)
		defer cancel()
	}

	// 执行任务
	startTime := time.Now()

	// 使用channel等待任务完成或超时
	done := make(chan error, 1)
	go func() {
		defer func() {
			if r := recover(); r != nil {
				done <- fmt.Errorf("task panicked: %v", r)
			}
		}()
		done <- task.Fn(ctx)
	}()

	select {
	case err := <-done:
		duration := time.Since(startTime)
		if err != nil {
			p.failed.Add(1)
			// 减少日志输出
			if p.failed.Load()%50 == 0 { // 每50个失败任务输出一次
				fmt.Printf("[COMPLETED] Task '%s' finished with error in %v: %v\n",
					task.Name, duration, err)
			}
		} else {
			p.success.Add(1)
			// 减少成功日志输出频率
			if p.success.Load()%100 == 0 { // 每100个成功任务输出一次
				fmt.Printf("[COMPLETED] Task '%s' finished successfully in %v\n",
					task.Name, duration)
			}
		}
	case <-ctx.Done():
		duration := time.Since(startTime)
		p.failed.Add(1)
		// 减少超时日志输出
		if p.failed.Load()%50 == 0 {
			fmt.Printf("[TIMEOUT] Task '%s' timed out after %v\n", task.Name, duration)
		}
	}
}

// monitor 监控协程数量
func (p *Pool) monitor() {
	ticker := time.NewTicker(3 * time.Second)
	defer ticker.Stop()
	defer close(p.monitorStopped)

	fmt.Println("[MONITOR] Started monitoring goroutines...")

	for {
		select {
		case <-ticker.C:
			running := p.GetRunningCount()
			// 只有在有运行中的协程时才显示统计信息
			if running > 0 {
				fmt.Println("\n=== 当前统计信息 ===")
				fmt.Printf("运行中协程数: %d\n", running)
				fmt.Printf("成功完成任务数: %d\n", p.GetSuccessCount())
				fmt.Printf("失败任务数: %d\n", p.GetFailedCount())
				fmt.Printf("总完成数: %d\n", p.GetTotalCompletedCount())

				// 显示每个CPU核心的协程数量
				if p.maxGoroutinesPerCore > 0 {
					fmt.Println("\n=== 各CPU核心协程分布 ===")
					for i := 0; i < p.numCPU; i++ {
						count := p.coreGoroutineCounts[i].Load()
						capacity := fmt.Sprintf("%d/%d", count, p.maxGoroutinesPerCore)
						fmt.Printf("CPU核心 %d: %s 个协程\n", i, capacity)
					}
				}
			}
		case <-p.stopMonitor:
			fmt.Println("[MONITOR] Stopped monitoring goroutines")
			return
		}
	}
}
